# 导入所需的库
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from torchvision import datasets, transforms
from lightning.fabric import Fabric, seed_everything
from torchmetrics.classification import Accuracy
import time
import argparse # 导入 argparse 用于接收命令行参数

# 定义简单的卷积神经网络模型 (保持不变)
class SimpleConvNet(nn.Module):
    def __init__(self):
        super(SimpleConvNet, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        # 假设输入 28x28 -> (28-3+1)/1 = 26 -> (26-3+1)/1 = 24 -> 24/2 = 12
        # 64 * 12 * 12 = 9216
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

# 训练一个 epoch 的函数 (保持不变)
def train_epoch(fabric: Fabric, model: nn.Module, train_loader: DataLoader, optimizer: optim.Optimizer, epoch: int):
    model.train()
    # enumerate(train_loader) 会在 DDP 下自动使用 DistributedSampler
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        fabric.backward(loss)
        optimizer.step()
        if batch_idx % 100 == 0 and fabric.is_global_zero: # 只在全局 rank 0 打印进度
            # 注意：len(data) 是单个设备的 batch size, 总处理数需要乘以 world_size
            processed_samples = batch_idx * len(data) * fabric.world_size + len(data) * (fabric.global_rank + 1)
            total_samples = len(train_loader.dataset)
            percent = 100. * (batch_idx * fabric.world_size * len(data)) / total_samples # 估算进度
            fabric.print(f'训练 Epoch: {epoch} [{processed_samples}/{total_samples} '
                         f'({percent:.0f}%)]\t损失: {loss.item():.6f}')

# 测试一个 epoch 的函数 (保持不变)
def test_epoch(fabric: Fabric, model: nn.Module, test_loader: DataLoader):
    model.eval()
    test_acc = Accuracy(task="multiclass", num_classes=10).to(fabric.device)
    total_loss_tensor = torch.tensor(0.0, device=fabric.device)
    with torch.no_grad():
        for data, target in test_loader:
            output = model(data)
            batch_loss = F.nll_loss(output, target, reduction='sum')
            total_loss_tensor += batch_loss
            test_acc.update(output, target) # 在每个 rank 上独立更新

    # 聚合所有进程的损失和样本数
    gathered_losses = fabric.all_gather(total_loss_tensor) # 收集所有 rank 的 sum loss
    # 在 fabric >= 2.1 中，可以直接聚合 metric 对象
    # gathered_acc = fabric.all_gather(test_acc) # 收集 metric 状态 (或直接用 compute + all_reduce)

    # 计算全局平均损失和准确率 (仅在 rank 0 计算和打印最终结果)
    if fabric.is_global_zero:
        avg_loss = gathered_losses.sum() / len(test_loader.dataset)
        # 对于 Accuracy metric，compute() 应该在所有进程调用后聚合 (或者用 fabric.all_reduce 手动聚合)
        # fabric.barrier() # 确保所有进程都完成了 update
        final_acc = test_acc.compute() # compute 应该内部处理了同步 (在较新版本 TorchMetrics/Fabric 中)
        total_samples = len(test_loader.dataset)
        correct_samples = final_acc * total_samples
        fabric.print(f'\n测试集: 平均损失: {avg_loss:.4f}, '
                     f'准确率: {final_acc*100:.0f}% ({int(correct_samples)}/{total_samples})\n')
    else:
        # 其他 rank 也需要调用 compute 来参与同步（取决于 metric 实现）
        test_acc.compute()

    # 所有 rank 都需要重置 metric
    test_acc.reset()
    fabric.barrier() # 确保所有进程在继续之前都已完成测试和打印


# 主执行函数
def main(args):
    # --- 超参数设置 ---
    batch_size_per_device = 64
    epochs = 5
    lr = 1.0
    seed = 42
    num_workers_per_loader = 2

    # --- Fabric 初始化 ---
    # **关键修改**: 添加 num_nodes 参数
    # `devices` 现在指每个节点上的设备数量
    fabric = Fabric(accelerator="cuda",
                    devices=args.gpus_per_node,  # 每个节点上的 GPU 数
                    num_nodes=args.num_nodes,    # 总节点数
                    strategy="ddp",
                    precision='16-mixed')
    # **关键**: fabric.launch() 不再负责启动进程，它假设进程已由外部启动器（如 torchrun）启动
    # 它在这里主要用于设置当前进程的环境（设备放置、seed 等）
    fabric.launch()

    # --- 可复现性 ---
    # 使用 Fabric 的 seed_everything 确保所有进程使用不同的、确定性的种子
    seed_everything(seed + fabric.global_rank)

    # --- 数据准备 ---
    # **关键**: 确保所有节点都能访问此路径！最好是共享文件系统。
    # 如果不是共享文件系统，需要确保在运行前，数据已存在于每个节点的这个路径下。
    # data_path = "/path/to/shared/mnist_data" # 示例：使用共享路径
    data_path = "./mnist_data" # 使用相对路径，假设所有节点运行目录结构相同
    fabric.print(f"全局 Rank {fabric.global_rank}: 使用数据路径: {data_path}")

    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])

    # --- 数据下载 (只在全局 Rank 0 执行) ---
    if fabric.is_global_zero:
        fabric.print(f"全局 Rank {fabric.global_rank}: 检查并下载 MNIST 数据集到 {data_path}...")
        # 确保目录存在
        os.makedirs(data_path, exist_ok=True)
        datasets.MNIST(data_path, train=True, download=True)
        datasets.MNIST(data_path, train=False, download=True)
        fabric.print(f"全局 Rank {fabric.global_rank}: 数据集下载完成或已存在。")

    # --- 同步点 ---
    # **关键**: 确保其他 rank 等待 rank 0 下载完成数据
    fabric.barrier()
    fabric.print(f"全局 Rank {fabric.global_rank}: 继续加载数据集...")

    # --- 加载数据集 (所有进程执行) ---
    train_dataset = datasets.MNIST(
        data_path, train=True, download=False, transform=transform # download=False 避免竞争
    )
    test_dataset = datasets.MNIST(
        data_path, train=False, download=False, transform=transform
    )

    # --- 创建数据加载器 ---
    # Fabric 的 setup_dataloaders 会自动处理 DistributedSampler
    train_loader = DataLoader(
        train_dataset,
        batch_size=batch_size_per_device,
        num_workers=num_workers_per_loader,
        # shuffle=True, # DDP 时不需要设置 shuffle=True，Sampler 会处理
        persistent_workers=True if num_workers_per_loader > 0 else False,
        pin_memory=True # 通常建议开启
    )
    test_loader = DataLoader(
        test_dataset,
        batch_size=batch_size_per_device * 2,
        num_workers=num_workers_per_loader,
        shuffle=False,
        persistent_workers=True if num_workers_per_loader > 0 else False,
        pin_memory=True
    )

    fabric.print(f"进程 {fabric.global_rank}/{fabric.world_size} (节点 Rank: {fabric.node_rank}, 本地 Rank: {fabric.local_rank}) 使用设备: {fabric.device}")
    fabric.print(f"训练数据集大小: {len(train_dataset)}")
    fabric.print(f"测试数据集大小: {len(test_dataset)}")
    fabric.print(f"每个设备的 Batch Size: {batch_size_per_device}")
    fabric.print(f"有效全局 Batch Size: {batch_size_per_device * fabric.world_size}")


    # --- 使用 Fabric 设置模型、优化器和数据加载器 ---
    # 注意: setup_dataloaders 会自动添加 DistributedSampler
    train_loader, test_loader = fabric.setup_dataloaders(train_loader, test_loader)

    # 实例化模型和优化器
    model = SimpleConvNet()
    # SyncBatchNorm: 如果模型包含 BatchNorm 层并且 world_size > 1, 推荐在 setup 之前转换
    if fabric.world_size > 1:
       fabric.print("模型包含 BatchNorm 且 world_size > 1, 尝试转换 SyncBatchNorm (如果模型有 BatchNorm 层)")
       # 注意: 只有当模型确实有 BatchNorm 层时才转换，否则可能出错或无效
       # 更好的做法是检查模型中是否有 BatchNorm 层
       has_batchnorm = any(isinstance(m, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)) for m in model.modules())
       if has_batchnorm:
           model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
           fabric.print("已转换 SyncBatchNorm")
       else:
           fabric.print("模型中未找到 BatchNorm 层，跳过 SyncBatchNorm 转换")

    optimizer = optim.Adadelta(model.parameters(), lr=lr)
    # setup 会处理模型移动到设备和 DDP 包装
    model, optimizer = fabric.setup(model, optimizer)

    # --- 训练循环 ---
    fabric.print(f"\n在 {fabric.strategy.num_nodes} 个节点, 共 {fabric.world_size} 个 GPU 上开始训练...")
    start_time = time.time()

    for epoch in range(1, epochs + 1):
        epoch_start_time = time.time()
        # 设置 DDP Sampler 的 epoch，确保每个 epoch 的 shuffle 不同
        # fabric.setup_dataloaders 内部可能会处理，但显式设置更安全
        if hasattr(train_loader.sampler, 'set_epoch'):
             train_loader.sampler.set_epoch(epoch)

        train_epoch(fabric, model, train_loader, optimizer, epoch)
        fabric.barrier() # 确保所有进程完成训练 epoch
        test_epoch(fabric, model, test_loader) # test_epoch 内部有 barrier
        # fabric.barrier() # test_epoch 结束时已有 barrier

        # 只在 rank 0 打印 epoch 时间
        if fabric.is_global_zero:
            epoch_time = time.time() - epoch_start_time
            fabric.print(f"Epoch {epoch} 在 {epoch_time:.2f} 秒内完成。")
        fabric.barrier() # 确保 rank 0 打印完再进入下一轮

    # 只在 rank 0 打印总时间
    if fabric.is_global_zero:
        total_time = time.time() - start_time
        fabric.print(f"\n训练在 {total_time:.2f} 秒内完成。")

    # --- 保存模型 (仅在全局 Rank 0 保存) ---
    save_path = "mnist_fabric_multinode_model.pt"
    state = {"model": model} # fabric.save 会自动处理提取 state_dict
    fabric.save(save_path, state)
    if fabric.is_global_zero:
        fabric.print(f"\n模型状态已保存到 {save_path} (由全局 Rank 0 保存)")

    fabric.barrier() # 确保保存完成后所有进程再退出
    fabric.print(f"全局 Rank {fabric.global_rank}: 训练和保存完成。")
    # fabric.launch() 的上下文管理器会自动处理退出时的清理

# 脚本入口点
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Lightning Fabric MNIST Multi-Node Example')
    # 添加多机多卡所需的参数
    parser.add_argument('--num_nodes', type=int, default=2, # 默认至少 2 个节点
                        help='Number of nodes for distributed training')
    parser.add_argument('--gpus_per_node', type=int, default=1, # 默认每个节点 1 个 GPU
                        help='Number of GPUs per node')
    # 你也可以添加其他参数，如 batch_size, epochs, lr, data_path 等
    # parser.add_argument('--data_path', type=str, default='./mnist_data', help='Path to dataset')
    args = parser.parse_args()

    # **重要**: CUDA 检查现在意义不大，因为脚本会在不同机器上运行
    # 启动器 (torchrun) 应该负责检查资源
    # 但可以保留一个本地检查作为参考
    if not torch.cuda.is_available():
        print("警告：当前环境未检测到 CUDA。假定目标节点有 CUDA 设备。")
    else:
        available_gpus = torch.cuda.device_count()
        print(f"本地检测到 {available_gpus} 个 GPU。脚本将根据 --gpus_per_node={args.gpus_per_node} 运行。")
        if available_gpus < args.gpus_per_node:
             print(f"警告：本地 GPU 数量 ({available_gpus}) 少于指定的 --gpus_per_node ({args.gpus_per_node})。确保目标节点有足够的 GPU。")

    main(args)